import pickle

import torch
from torch.utils.data import Dataset


class MyDataset(Dataset):
    def __init__(self, input_list, max_len):
        super(MyDataset, self).__init__()
        self.input_list = input_list
        self.max_len = max_len

    def __len__(self):
        return len(self.input_list)

    def __getitem__(self, item):
        input_ids = self.input_list[item]
        input_ids = input_ids[:self.max_len]
        input_ids = torch.tensor(input_ids, dtype=torch.long)
        return input_ids


if __name__ == '__main__':
    pkl_path = '../data/medical_valid.pkl'
    with open(pkl_path, 'rb') as f:
        input_list = pickle.load(f)
        print(len(input_list))
        print(input_list[-2])
    max_len = 300
    dataset = MyDataset(input_list, max_len)
    print(len(dataset))
    print(dataset[-2])
